import requests
import json
import time
import os
from datetime import datetime


def try_on_clothes(api_key, person_image_url, top_garment_url=None, bottom_garment_url=None):
    """
    调用阿里云DashScope虚拟试衣API

    参数:
        api_key: DashScope API密钥
        person_image_url: 人物图片URL
        top_garment_url: 上衣图片URL（可选）
        bottom_garment_url: 下装图片URL（可选）

    返回:
        试衣结果图片的URL
    """
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis/"

    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    # 构建请求数据
    data = {
        "model": "aitryon",
        "input": {
            "person_image_url": person_image_url
        },
        "parameters": {
            "resolution": -1,
            "restore_face": True
        }
    }

    # 添加可选的服装图片
    if top_garment_url:
        data["input"]["top_garment_url"] = top_garment_url
    if bottom_garment_url:
        data["input"]["bottom_garment_url"] = bottom_garment_url

    try:
        # 发送异步请求
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()

        print(f"API响应: {json.dumps(result, ensure_ascii=False, indent=2)}")

        # 获取任务ID
        task_id = result.get("output", {}).get("task_id")
        if not task_id:
            raise Exception("未获取到任务ID")

        print(f"任务ID: {task_id}")

        # 轮询获取结果
        max_retries = 30  # 最多等待60秒
        retry_count = 0

        while retry_count < max_retries:
            time.sleep(2)  # 等待2秒
            status_response = requests.get(
                f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}",
                headers={"Authorization": f"Bearer {api_key}"}
            )
            status_result = status_response.json()
            print(
                f"任务状态 ({retry_count + 1}/{max_retries}): {json.dumps(status_result, ensure_ascii=False, indent=2)}")

            task_status = status_result.get("output", {}).get("task_status")
            if task_status == "SUCCEEDED":
                # 修正：从正确的字段提取图片URL
                image_url = status_result.get("output", {}).get("image_url")
                if not image_url:
                    # 备选路径，尝试其他可能的字段名
                    image_url = status_result.get("output", {}).get("url")

                if not image_url:
                    raise Exception("无法从API响应中获取结果图片URL")

                return image_url
            elif task_status == "FAILED":
                error_msg = status_result.get("output", {}).get("error", {}).get("message", "未知错误")
                raise Exception(f"任务失败: {error_msg}")

            retry_count += 1

        raise Exception("任务超时")

    except requests.exceptions.RequestException as e:
        print(f"请求错误: {e}")
        raise
    except Exception as e:
        print(f"发生错误: {e}")
        raise


def download_image(url, save_path):
    """
    下载图片并保存到指定路径

    参数:
        url: 图片URL
        save_path: 保存路径

    返回:
        保存的文件路径
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()

        # 保存图片
        with open(save_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)

        print(f"图片已保存到: {save_path}")
        return save_path
    except Exception as e:
        print(f"下载图片失败: {e}")
        raise


def main():
    # 配置信息
    API_KEY = os.environ.get("ali_API_KEY")
    if not API_KEY:
        API_KEY = input("请输入阿里云DashScope API密钥: ")
        if not API_KEY:
            raise ValueError("API密钥不能为空")
        # 临时设置环境变量，仅本次运行有效
        os.environ["ali_API_KEY"] = API_KEY

    # 图片URL
    模特人物照片链接 = "https://help-static-aliyun-doc.aliyuncs.com/assets/img/zh-CN/1389646171/p801328.png"
    上身衣服照片链接 = "https://help-static-aliyun-doc.aliyuncs.com/assets/img/zh-CN/1389646171/p801339.jpeg"
    下身衣服照片链接 = "https://help-static-aliyun-doc.aliyuncs.com/assets/img/zh-CN/1389646171/p801340.jpeg"

    try:
        # 调用试衣API
        result_url = try_on_clothes(
            API_KEY,
            模特人物照片链接,
            上身衣服照片链接,
            下身衣服照片链接
        )

        print(f"试衣成功！结果图片URL: {result_url}")

        # 生成带时间戳的文件名，防止覆盖
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = f"tryon_result_{timestamp}.jpg"

        # 下载并保存图片
        downloaded_path = download_image(result_url, save_path)
        print(f"试衣图片已保存到当前目录: {os.path.abspath(downloaded_path)}")

    except Exception as e:
        print(f"试衣失败: {str(e)}")


if __name__ == '__main__':
    main()